module dycore_mod
  use gptl_timer_mod
  use physics_const_mod
  use spark_cubed_sphere
  use params_mod
  use time_mod
  use static_mod
  use state_mod
  use diag_mod
  use tend_mod
  use math_mod
  use riemann_mod
  use ghost_interp_mod
  use hope_recon_mod
  use test_mod
  use time_schemes_mod

  implicit none

  private

  public dycore_init
  public dycore_run
  public dycore_final
  public dycore_diag_fields
  public dycore_diag_scalars
  public static_type
  public static
  public state_type
  public tend_type
  
  interface
    subroutine dycore_diag_fields_interface(state, diag)
      import state_type, diag_type
      type(state_type), intent(inout) :: state
      type(diag_type ), intent(inout) :: diag
    end subroutine dycore_diag_fields_interface
    
    subroutine dycore_diag_scalars_interface(state, diag, static)
      import state_type, diag_type, static_type
      type(state_type ), intent(inout) :: state
      type(diag_type  ), intent(inout) :: diag
      type(static_type), intent(in   ) :: static
    end subroutine dycore_diag_scalars_interface
  end interface
  
  procedure(time_integrator_interface), pointer :: time_integrator => null()

  procedure(dycore_diag_fields_interface), pointer :: dycore_diag_fields => null()

  procedure(dycore_diag_scalars_interface), pointer :: dycore_diag_scalars => null()

contains

  subroutine dycore_init(mesh)

    type(cubed_sphere_mesh_type), intent(in) :: mesh

    integer i

    call dycore_final()
    
    call diag%init(mesh)
    call static%init(mesh)
    allocate(state(nsubstep))
    allocate(tend (nsubstep))
    
    do i = 1, nsubstep
      call state(i)%init(mesh)
      call tend (i)%init(mesh)
    end do
    
    call state_ref%init(mesh)
    
    if( trim(time_scheme) == 'IMEX' )then
      allocate( state_IMEX(IRK_stage+1) )
      allocate( tend_IM   (IRK_stage+1) )
      allocate( tend_EX   (IRK_stage+1) )
      do i = 1, IRK_stage+1
        call state_IMEX(i)%init(mesh)
        call tend_IM   (i)%init(mesh)
        call tend_EX   (i)%init(mesh)
      enddo
      call state_IRK%init(mesh)
      call state_IM %init(mesh)
      call tend_IRK %init(mesh)
      
      allocate( state_HON(4) ) ! Adjust for specified Newton method
      allocate( tend_HON (4) ) ! Adjust for specified Newton method
      do i = 1,4
        call state_HON(i)%init(mesh)
        call tend_HON (i)%init(mesh)
      enddo
    endif
    
    select case(model_type)
    case(0)
      dycore_diag_fields  => dycore_diag_fields_adv
      dycore_diag_scalars => dycore_diag_scalars_adv
    case(1)
      dycore_diag_fields  => dycore_diag_fields_barotropic
      dycore_diag_scalars => dycore_diag_scalars_barotropic
    case(2)
      dycore_diag_fields  => dycore_diag_fields_baroclinic
      dycore_diag_scalars => dycore_diag_scalars_baroclinic
    end select

    select case (time_scheme)
    case ('Euler')
      time_integrator => euler
    case ('SSP_RK3')
      time_integrator => ssp_rk3
    case ('WRF_RK3')
      time_integrator => WRF_RK3
    case ('RK4')
      time_integrator => rk4
    case ('IMEX')
      time_integrator => IMEX
    end select

    call riemann_init()
    call ghost_interp_init(mesh)
    call hope_recon_init(mesh)
    
  end subroutine dycore_init

  subroutine dycore_run(old, new, dt)

    integer, intent(in) :: old
    integer, intent(in) :: new
    real(8), intent(in) :: dt

    call time_integrator(old, new, dt, state, tend)

  end subroutine dycore_run

  subroutine dycore_final()

    integer i

    call static%clear()
    if (allocated(state)) then
      do i = 1, nsubstep
        call state(i)%clear()
        call tend (i)%clear()
      end do
      deallocate(state)
      deallocate(tend)
    end if

    call ghost_interp_final()
    call hope_recon_final()

  end subroutine dycore_final
  
  subroutine dycore_diag_fields_adv(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    real(r8) uc, vc, u, v
    real(r8) JhQ(ncq)
    real(r8), dimension(2,2) :: Ah
    integer ids, ide, jds, jde, kds, kde, pc, pqs, pqe
    integer idom, p, ip, i, j, k
    
    ret = gptlstart('dycore_diag_fields_adv')

    do idom = 1, mesh%ndom
      if (state%mesh%domain(idom)%active) then
        call state%mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, &
                                   kds=kds, kde=kde, pc=pc, pqs=pqs, pqe=pqe)
        associate (domain => state%mesh%domain(idom)   , &
                   dx     => state%mesh%domain(idom)%dx, &
                   dy     => state%mesh%domain(idom)%dy, &
                   q      => state%domain(idom)%q      , &
                   gd     => diag%domain(idom)%gd      , &
                   ght    => diag%domain(idom)%ght     , &
                   us     => diag%domain(idom)%us      , &
                   vs     => diag%domain(idom)%vs      , &
                   u      => diag%domain(idom)%u       , &
                   v      => diag%domain(idom)%v       , &
                   uu     => diag%domain(idom)%uc      , &
                   vv     => diag%domain(idom)%vc      , &
                   vor    => diag%domain(idom)%vor     , &
                   ghs    => static%domain(idom)%ghs)
        !$OMP PARALLEL DO PRIVATE(i,j,ip,p,JhQ,uc,vc,Ah) COLLAPSE(3)
        do k = kds, kde
          do j = jds - 1, jde + 1
            do i = ids - 1, ide + 1
              do ip = 1, ncq
                p = pqs + ip - 1
                JhQ(ip) = mesh%get_dom_Jh(idom,p,i,j,k)
              enddo
              gd (i,j,k) = q(i,j,k,1) / domain%cell_quad(JhQ)
              ght(i,j,k) = gd(i,j,k) + ghs(i,j)
              uc = uu(i,j,k)
              vc = vv(i,j,k)
              Ah = mesh%get_dom_Ah(idom,pc,i,j,k)
              call contrav_to_spherev(Ah, uc, vc, us(i,j,k), vs(i,j,k))
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    
    ret = gptlstop('dycore_diag_fields_adv')

  end subroutine dycore_diag_fields_adv

  subroutine dycore_diag_fields_barotropic(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    real(r8) uc, vc, u, v
    real(r8), dimension(2,2) :: Gh, Ah
    real(r8) JhQ(ncq)
    integer ids, ide, jds, jde, kds, kde, pc, pqs, pqe
    integer idom, ip, p, i, j, k
    
    ret = gptlstart('dycore_diag_fields_barotropic')

    do idom = 1, mesh%ndom
      if (state%mesh%domain(idom)%active) then
        call state%mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, &
                                   kds=kds, kde=kde, pc=pc, pqs=pqs, pqe=pqe)
        associate (domain => state%mesh%domain(idom)   , &
                   dx     => state%mesh%domain(idom)%dx, &
                   dy     => state%mesh%domain(idom)%dy, &
                   q      => state%domain(idom)%q      , &
                   gd     => diag%domain(idom)%gd      , &
                   ght    => diag%domain(idom)%ght     , &
                   us     => diag%domain(idom)%us      , &
                   vs     => diag%domain(idom)%vs      , &
                   u      => diag%domain(idom)%u       , &
                   v      => diag%domain(idom)%v       , &
                   vor    => diag%domain(idom)%vor     , &
                   ghs    => static%domain(idom)%ghs)
        do k = kds, kde
          do j = jds - 1, jde + 1
            do i = ids - 1, ide + 1
              do ip = 1, ncq
                p = pqs + ip - 1
                JhQ(ip) = mesh%get_dom_Jh(idom,p,i,j,k)
              enddo
              gd (i,j,k) = q(i,j,k,1) / domain%cell_quad(JhQ)
              ght(i,j,k) = gd(i,j,k) + ghs(i,j)
              uc = q(i,j,k,2) / q(i,j,k,1)
              vc = q(i,j,k,3) / q(i,j,k,1)
              Gh = mesh%get_dom_Gh(idom,pc,i,j,k)
              Ah = mesh%get_dom_Ah(idom,pc,i,j,k)
              call contrav_to_spherev(Ah, uc, vc, us(i,j,k), vs(i,j,k))
              call contrav_to_cov    (Gh, uc, vc, u (i,j,k), v (i,j,k))
            end do
          end do
        end do
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide
              vor(i,j,k) = (center_diff2(v(i-1:i+1,j,k), dx) - center_diff2(u(i,j-1:j+1,k), dy)) / mesh%get_dom_Jh(idom,pc,i,j,k)
            end do
          end do
        end do
        end associate
      end if
    end do
    
    ret = gptlstop('dycore_diag_fields_barotropic')

  end subroutine dycore_diag_fields_barotropic

  subroutine dycore_diag_fields_baroclinic(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    real(r8) uc, vc
    real(r8), dimension(2,2) :: Gh, Ah
    integer idom, ids, ide, jds, jde, kds, kde, pc, pqs, pqe, i, j, k
    
    ret = gptlstart('dycore_diag_fields_baroclinic')

    do idom = 1, mesh%ndom
      if (state%mesh%domain(idom)%active) then
        call state%mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde,&
                                   pc=pc, pqs=pqs, pqe=pqe)
        associate (domain => state%mesh%domain(idom),    &
                   dx     => state%mesh%domain(idom)%dx, &
                   dy     => state%mesh%domain(idom)%dy, &
                   q      => state%domain(idom)%q,       &
                   us     => diag%domain(idom)%us,       &
                   vs     => diag%domain(idom)%vs,       &
                   u      => diag%domain(idom)%u,        &
                   v      => diag%domain(idom)%v,        &
                   vor    => diag%domain(idom)%vor)
        !$OMP PARALLEL
        !$OMP DO PRIVATE(i,j,uc,vc,Gh,Ah) COLLAPSE(3)
        do k = kds, kde
          do j = jds - 1, jde + 1
            do i = ids - 1, ide + 1
              uc = q(i,j,k,2) / ( q(i,j,k,1) + q(i,j,k,6) )
              vc = q(i,j,k,3) / ( q(i,j,k,1) + q(i,j,k,6) )
              Gh = mesh%get_dom_Gh(idom,pc,i,j,k)
              Ah = mesh%get_dom_Ah(idom,pc,i,j,k)
              call contrav_to_spherev(Ah, uc, vc, us(i,j,k), vs(i,j,k))
              call contrav_to_cov    (Gh, uc, vc, u (i,j,k), v (i,j,k))
            end do
          end do
        end do
        !$OMP END DO
        !$OMP DO PRIVATE(i,j) COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide
              vor(i,j,k) = (center_diff2(v(i-1:i+1,j,k), dx) - center_diff2(u(i,j-1:j+1,k), dy)) / mesh%get_dom_Jh(idom,pc,i,j,k)
            end do
          end do
        end do
        !$OMP END DO
        !$OMP END PARALLEL
        end associate
      end if
    end do

    ret = gptlstop('dycore_diag_fields_baroclinic')
    
  end subroutine dycore_diag_fields_baroclinic

  subroutine dycore_diag_scalars_adv(state, diag, static)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    type(static_type), intent(in) :: static

    integer idom, ids, ide, jds, jde, i, j
    real(r8) uc, vc, u, v, ke, pe
    real(r8) :: mass, energy

    logical , save :: first_call = .true.
    real(r8), save :: tm0, te0
    
    ret = gptlstart('dycore_diag_scalars_adv')

    mass = 0
    energy = 0
    associate(mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde)
        associate (dx  => mesh%domain(idom)%dx,  &
                   dy  => mesh%domain(idom)%dy,  &
                   q   => state%domain(idom)%q,  &
                   us  => diag%domain(idom)%us,  &
                   vs  => diag%domain(idom)%vs,  &
                   ghs => static%domain(idom)%ghs)
        !$OMP PARALLEL DO PRIVATE(pe,i) COLLAPSE(2) REDUCTION(+:mass,energy)
        do j = jds, jde
          do i = ids, ide
            mass = mass + state%domain(idom)%q(i,j,1,1) * dx * dy
            pe = 0.5 * q(i,j,1,1)**2
            energy = energy + pe * dx * dy
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

    diag%tm = global_sum(mass)
    diag%te = global_sum(energy)

    if (first_call) then
      tm0 = diag%tm
      te0 = diag%te
      first_call = .false.
    end if

    call log_add_diag('tm', (diag%tm - tm0) / tm0)
    call log_add_diag('te', (diag%te - te0) / te0)
    
    ret = gptlstop('dycore_diag_scalars_adv')

  end subroutine dycore_diag_scalars_adv
  
  subroutine dycore_diag_scalars_barotropic(state, diag, static)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    type(static_type), intent(in) :: static

    integer idom, ids, ide, jds, jde, i, j
    real(r8) uc, vc, u, v, ke, pe
    real(r8) Jh
    real(r8) Gh(2,2)
    real(r8) :: mass, energy

    logical , save :: first_call = .true.
    real(r8), save :: tm0, te0
    
    ret = gptlstart('dycore_diag_scalars_barotropic')

    mass = 0
    energy = 0
    associate(mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde)
        associate (dx  => mesh%domain(idom)%dx,  &
                   dy  => mesh%domain(idom)%dy,  &
                   q   => state%domain(idom)%q,  &
                   us  => diag%domain(idom)%us,  &
                   vs  => diag%domain(idom)%vs,  &
                   ghs => static%domain(idom)%ghs)
        do j = jds, jde
          do i = ids, ide
            mass = mass + state%domain(idom)%q(i,j,1,1) * dx * dy
            uc = q(i,j,1,2) / q(i,j,1,1)
            vc = q(i,j,1,3) / q(i,j,1,1)
            Gh = mesh%get_dom_Gh(idom,1,i,j,1)
            Jh = mesh%get_dom_Jh(idom,1,i,j,1)
            call contrav_to_cov(Gh, uc, vc, u, v)
            ke = 0.5 * (q(i,j,1,2) * u + q(i,j,1,3) * v)
            pe = 0.5 * (q(i,j,1,1) + Jh * ghs(i,j))**2
            energy = energy + (ke + pe) * dx * dy
          end do
        end do
        end associate
      end if
    end do
    end associate

    diag%tm = global_sum(mass)
    diag%te = global_sum(energy)

    if (first_call) then
      tm0 = diag%tm
      te0 = diag%te
      first_call = .false.
    end if

    call log_add_diag('tm', (diag%tm - tm0) / tm0)
    call log_add_diag('te', (diag%te - te0) / te0)
    
    ret = gptlstop('dycore_diag_scalars_barotropic')

  end subroutine dycore_diag_scalars_barotropic

  subroutine dycore_diag_scalars_baroclinic(state, diag, static)
    type(state_type ), intent(inout) :: state
    type(diag_type  ), intent(inout) :: diag
    type(static_type), intent(in   ) :: static

    integer idom, ids, ide, jds, jde, kds, kde
    integer pc, pqs, pqe
    integer i, j, k
    real(r8) uc, vc, u, v
    real(r8) ke, pe, ie
    real(r8) sqrtG, alpha, mc, Cp, Cv
    real(r8) Gh(2,2)
    logical , save :: first_call = .true.
    real(r8), save :: tm0, te0
    real(r8) :: mass, energy
    
    ret = gptlstart('dycore_diag_scalars_baroclinic')

    mass = 0
    energy = 0
    associate(mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pc=pc, pqs=pqs, pqe=pqe)
        associate (Jab    => mesh%domain(idom)%Jab,     &
                   dx     => mesh%domain(idom)%dx,      &
                   dy     => mesh%domain(idom)%dy,      &
                   deta   => mesh%domain(idom)%deta,    &
                   z      => mesh%domain(idom)%z,       &
                   q      => state%domain(idom)%q,      &
                   rhod   => diag%domain(idom)%rhod,    &
                   rhom   => diag%domain(idom)%rhom,    &
                   rho    => diag%domain(idom)%rho,     &
                   p      => diag%domain(idom)%p,       &
                   pt     => diag%domain(idom)%pt,      &
                   t      => diag%domain(idom)%t,       &
                   mr     => diag%domain(idom)%mr,      &
                   w      => diag%domain(idom)%w,       &
                   g      => static%domain(idom)%gravityQ)
          !$OMP PARALLEL DO PRIVATE(sqrtG,alpha,mc,Cp,Cv,ke,pe,ie,uc,vc,u,v,Gh) COLLAPSE(3) REDUCTION(+:mass,energy)
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                sqrtG = Jab(pc,i,j,k)
                rhod(i,j,k) = q(i,j,k,1)
                rhom(i,j,k) = q(i,j,k,6)
                rho (i,j,k) = rhod(i,j,k) + rhom(i,j,k)
                mr  (i,j,k) = q(i,j,k,6) / q(i,j,k,1)
                pt  (i,j,k) = q(i,j,k,5) / q(i,j,k,1)
                call calc_thermal_parameters(mr(i,j,k), Cp=Cp, Cv=Cv, mc=mc, alpha=alpha)
                p   (i,j,k) = p0 * ( rho(i,j,k) * Rd * mc * pt(i,j,k) / p0 )**(Cp / Cv)
                t   (i,j,k) = p(i,j,k) / ( rho(i,j,k) * Rd * mc )
                w   (i,j,k) = q(i,j,k,4) / ( q(i,j,k,1) + q(i,j,k,6) )
                
                uc = q(i,j,k,2) / ( q(i,j,k,1) + q(i,j,k,6) )
                vc = q(i,j,k,3) / ( q(i,j,k,1) + q(i,j,k,6) )
                Gh = mesh%get_dom_Gh(idom,pc,i,j,k)
                call contrav_to_cov(Gh, uc, vc, u, v)
                ke = 0.5 * (q(i,j,k,2) * u + q(i,j,k,3) * v + q(i,j,k,4) * w(i,j,k))
                pe = rho(i,j,k) * mesh%domain(idom)%cell_quad( z(pqs:pqe,i,j,k) * g(pqs:pqe,i,j,k) )
                ie = rho(i,j,k) * Cp * t(i,j,k)
                
                mass   = mass + rho (i,j,k) * Jab(pc,i,j,k) * dx * dy * deta
                energy = energy + ( ke + pe + ie ) * Jab(pc,i,j,k) * dx * dy * deta
              end do
            end do
          end do
          !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

    diag%tm = global_sum(mass)
    diag%te = global_sum(energy)
    
    if (first_call) then
      tm0 = diag%tm
      te0 = diag%te
      first_call = .false.
    end if

    call log_add_diag('tm', (diag%tm - tm0) / tm0)
    call log_add_diag('te', (diag%te - te0) / te0)
    
    ret = gptlstop('dycore_diag_scalars_baroclinic')

  end subroutine dycore_diag_scalars_baroclinic
  
end module dycore_mod
